from abc import ABC, abstractmethod
from sklearn.model_selection import LeaveOneOut, GridSearchCV
from sklearn.metrics import accuracy_score, recall_score, confusion_matrix, roc_curve, auc
from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier
from sklearn.neural_network import MLPClassifier
from sklearn.ensemble import RandomForestClassifier
from .dataset import DataPreprocessor
import numpy as np


# ==================== 基类定义 ====================
class BaseModel(ABC):
    """抽象基类，定义接口"""

    def __init__(self, preprocessor: DataPreprocessor):
        self.preprocessor = preprocessor
        self.best_model = None
        self.best_params_ = None
        self.y_true = []
        self.y_pred = []
        self.y_prob = []

    @abstractmethod
    def _get_param_grid(self):
        """参数网格，用于网格化搜索"""
        pass

    @abstractmethod
    def _create_model(self):
        """获取模型，子类必须实现"""
        pass

    def _prepare_data(self, x, y=None):
        if y is None:
            return self.preprocessor.preprocess(x)
        return self.preprocessor.preprocess(x, y)

    def train(self, x, y):
        """训练流程"""
        x_preprocessed, y_preprocessed = self._prepare_data(x, y)

        param_grid = self._get_param_grid()
        grid_search = GridSearchCV(
            estimator=self._create_model(),
            param_grid=param_grid,
            cv=LeaveOneOut(),
            scoring='accuracy',
            n_jobs=-1
        )
        grid_search.fit(x_preprocessed, y_preprocessed)
        self.best_model = grid_search.best_estimator_
        self.best_params_ = grid_search.best_params_

    def evaluate(self, x, y):
        """执行评估流程"""
        x_preprocessed = self._prepare_data(x)
        self.y_true.extend(y)
        self.y_pred.extend(self.best_model.predict(x_preprocessed))
        self.y_prob.extend(self.best_model.predict_proba(x_preprocessed)[:, 1])

    def get_metrics(self):
        """计算评估指标"""
        auc_value, fpr, tpr = self._calculate_auc()
        return {
            "accuracy": accuracy_score(self.y_true, self.y_pred),
            "sensitivity": recall_score(self.y_true, self.y_pred),
            "specificity": confusion_matrix(self.y_true, self.y_pred)[0, 0] /
                           (confusion_matrix(self.y_true, self.y_pred)[0, 0] +
                            confusion_matrix(self.y_true, self.y_pred)[0, 1]),
            "auc": auc_value,
            "fpr": fpr,
            "tpr": tpr
        }

    def _calculate_auc(self):
        """计算AUC值"""
        fpr, tpr, _ = roc_curve(self.y_true, self.y_prob)
        return auc(fpr, tpr), fpr, tpr


# ==================== 具体模型类 ====================
class SVMModel(BaseModel):
    """SVM模型实现类"""

    def _create_model(self):
        return SVC(probability=True, random_state=42)

    def _get_param_grid(self):
        return {
            'C': [0.1, 1, 10],
            'gamma': [0.01, 0.1, 0.001, 'scale', 'auto'],
            'kernel': ['rbf', 'linear', 'poly']
        }


class DecisionTreeModel(BaseModel):
    """决策树模型实现类"""

    def _create_model(self):
        return DecisionTreeClassifier(random_state=42)

    def _get_param_grid(self):
        return {
            'criterion': ['gini', 'entropy'],
            'max_depth': [3, 5, 7, 10, None],
            'min_samples_leaf': [1, 2, 4],
            'min_samples_split': [2, 5, 10]
        }


class NeuralNetworkModel(BaseModel):
    """神经网络模型实现类"""

    def _create_model(self):
        return MLPClassifier(max_iter=3000, random_state=42)

    def _get_param_grid(self):
        return {
            'hidden_layer_sizes': [(50,), (100,), (150)],
            'activation': ['relu', 'tanh', 'logistic'],
            'solver': ['adam', 'sgd'],
            'alpha': [0.0001, 0.001, 0.01]
        }


class RandomForest(BaseModel):
    def _create_model(self):
        return RandomForestClassifier(random_state=42)

    def _get_param_grid(self):
        return {
            "n_estimators": [50,100,200],
            "max_features": ["auto", "sqrt", "log2"],
            "max_depth": [10,20,None],
            "min_samples_leaf": [1, 2, 4]
        }
